package com.whz.gateway.common.web;

import com.google.common.collect.Sets;
import com.whz.gateway.common.convert.Convert;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.factory.rewrite.GzipMessageBodyResolver;
import org.springframework.cloud.gateway.filter.factory.rewrite.MessageBodyDecoder;
import org.springframework.cloud.gateway.filter.factory.rewrite.MessageBodyEncoder;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.ReactiveHttpOutputMessage;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.lang.NonNull;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.util.List;
import java.util.Map;
import java.util.function.Predicate;
import java.util.stream.Collectors;

/**
 * @author whz
 */
class LocalServerHttpResponseDecorator<T, R> extends ServerHttpResponseDecorator {

    final Map<String, MessageBodyDecoder> messageBodyDecoders;
    final Map<String, MessageBodyEncoder> messageBodyEncoders;
    private final Predicate<ServerWebExchange> predicate;
    private final ServerWebExchange exchange;
    private final Convert<T, R> convert;

    LocalServerHttpResponseDecorator(Predicate<ServerWebExchange> predicate,
                                     ServerWebExchange exchange,
                                     Convert<T, R> convert) {
        super(exchange.getResponse());
        this.predicate = predicate;
        this.exchange = exchange;
        this.convert = convert;
        messageBodyDecoders = Sets.newHashSet(new GzipMessageBodyResolver()).stream()
                                  .collect(Collectors.toMap(MessageBodyDecoder::encodingType, p -> p));
        messageBodyEncoders = Sets.newHashSet(new GzipMessageBodyResolver()).stream()
                                  .collect(Collectors.toMap(MessageBodyEncoder::encodingType, p -> p));
    }

    @Override
    @NonNull
    public Mono<Void> writeWith(@NonNull Publisher<? extends DataBuffer> body) {
        // 获取ContentType，判断是否返回JSON格式数据
        if (predicate.test(exchange)) {
            ClientResponse clientResponse = WebBodyUtils.getCurrentClientResponse(exchange, body);

            // 注意解码
            Class<T> inClass = convert.getInClass();
            Mono<R> modifiedBody = extractBody(exchange, clientResponse, inClass)
                    .flatMap(currentData -> Mono.defer(
                            () -> Mono.from(convert.getRewriteFunction().apply(exchange, currentData))))
                    .switchIfEmpty(
                            Mono.defer(() -> Mono.from(convert.getRewriteFunction().apply(exchange, null))));

            Class<R> outClass = convert.getOutClass();
            BodyInserter<Mono<R>, ReactiveHttpOutputMessage> bodyInserter =
                    BodyInserters.fromPublisher(modifiedBody, outClass);
            CachedHttpOutputMessage outputMessage =
                    new CachedHttpOutputMessage(exchange, exchange.getResponse().getHeaders());
            return bodyInserter.insert(outputMessage, new BodyInserterContext())
                               .then(Mono.defer(() -> {
                                   Mono<DataBuffer> messageBody =
                                           writeBody(getDelegate(), outputMessage, outClass);
                                   HttpHeaders headers = getDelegate().getHeaders();
                                   if (!headers.containsKey(HttpHeaders.TRANSFER_ENCODING) ||
                                       headers.containsKey(HttpHeaders.CONTENT_LENGTH)) {
                                       messageBody = messageBody.doOnNext(
                                               data -> headers.setContentLength(data.readableByteCount()));
                                   }

                                   return getDelegate().writeWith(messageBody);
                               }));
        }
        return super.writeWith(body);
    }

    @Override
    public @NonNull
    Mono<Void> writeAndFlushWith(@NonNull Publisher<? extends Publisher<? extends DataBuffer>> body) {
        return writeWith(Flux.from(body).flatMapSequential(p -> p));
    }

    private Mono<T> extractBody(ServerWebExchange exchange, ClientResponse clientResponse, Class<T> inClass) {
        // if inClass is byte[] then just return body, otherwise check if
        // decoding required
        if (byte[].class.isAssignableFrom(inClass)) {
            return clientResponse.bodyToMono(inClass);
        }

        List<String> encodingHeaders = exchange.getResponse().getHeaders()
                                               .getOrEmpty(HttpHeaders.CONTENT_ENCODING);
        for (String encoding : encodingHeaders) {
            MessageBodyDecoder decoder = messageBodyDecoders.get(encoding);
            if (decoder != null) {
                return clientResponse.bodyToMono(byte[].class)
                                     .publishOn(Schedulers.parallel())
                                     .map(decoder::decode)
                                     .map(bytes -> {
                                         final DataBuffer buffer =
                                                 exchange.getResponse().bufferFactory().wrap(bytes);
                                         return WebBodyUtils.getCurrentClientResponse(exchange, Mono.just(buffer));
                                     })
                                     .flatMap(response -> response.bodyToMono(inClass));
            }
        }

        return clientResponse.bodyToMono(inClass);
    }

    private Mono<DataBuffer> writeBody(ServerHttpResponse httpResponse, CachedHttpOutputMessage message,
                                       Class<?> outClass) {
        Mono<DataBuffer> response = DataBufferUtils.join(message.getBody());
        if (byte[].class.isAssignableFrom(outClass)) {
            return response;
        }

        List<String> encodingHeaders = httpResponse.getHeaders().getOrEmpty(HttpHeaders.CONTENT_ENCODING);
        for (String encoding : encodingHeaders) {
            MessageBodyEncoder encoder = messageBodyEncoders.get(encoding);
            if (encoder != null) {
                DataBufferFactory dataBufferFactory = httpResponse.bufferFactory();
                response = response.publishOn(Schedulers.parallel()).map(buffer -> {
                    byte[] encodedResponse = encoder.encode(buffer);
                    DataBufferUtils.release(buffer);
                    return encodedResponse;
                }).map(dataBufferFactory::wrap);
                break;
            }
        }

        return response;
    }
}